﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using NPredictor.Utils;

namespace NPredictor.DictionaryPredictor
{
	public class DictionaryPredictor<TWeight> : IWordPredictor<TWeight>
	{
		private readonly IDictionary<String, TWeight> wordsPool = new Dictionary<string, TWeight>();
		private readonly IDictionary<String, SortedDictionary<TWeight, IList<String>>> preffixList = new Dictionary<string, SortedDictionary<TWeight, IList<string>>>();
		private readonly IComparer<TWeight> comparer;

		public DictionaryPredictor(Func<TWeight, TWeight, int> comparer)
		{
			this.comparer = new WeightComparer<TWeight>(comparer);
		}

		public void Insert(string word, TWeight weight)
		{
			bool wordAlreadyExists;
			TWeight oldWeight = default(TWeight);

			if (wordAlreadyExists = wordsPool.ContainsKey(word))
				oldWeight = wordsPool[word];

			wordsPool[word] = weight;

			foreach (var preffix in word.Prefixes())
			{
				if (wordAlreadyExists)
				{
					RemoveWordFromPreffixList(preffix, word, oldWeight);
				}
				else if (!preffixList.ContainsKey(preffix))
				{
					preffixList[preffix] = new SortedDictionary<TWeight, IList<string>>(comparer);
				}

				if (!preffixList[preffix].ContainsKey(weight))
				{
					preffixList[preffix].Add(weight, new List<string>());
				}

				preffixList[preffix][weight].Add(word);
			}
		}

		public void Remove(string word)
		{
			TWeight weight = wordsPool[word];

			wordsPool.Remove(word);

			foreach (var preffix in word.Prefixes())
				RemoveWordFromPreffixList(preffix, word, weight);

		}

		public bool Search(string word, out TWeight weight)
		{
			if (wordsPool.ContainsKey(word))
			{
				weight = wordsPool[word];

				return true;
			}

			weight = default(TWeight);

			return false;
		}

		public IEnumerator<PredictionResult<TWeight>> Predict(string preffix)
		{
			if (preffixList.ContainsKey(preffix))
			{
				foreach (var result in preffixList[preffix])
					foreach (var word in result.Value)
						yield return new PredictionResult<TWeight>(word, result.Key);
			}
		}

		private void RemoveWordFromPreffixList(string preffix, string word, TWeight weight)
		{
			preffixList[preffix][weight].Remove(word);

			if (preffixList[preffix][weight].Count == 0)
				preffixList[preffix].Remove(weight);
		}
	}

	class WeightComparer<TWeight> : IComparer<TWeight>
	{
		private readonly Func<TWeight, TWeight, int> comparer;

		public int Compare(TWeight x, TWeight y)
		{
			return comparer(x, y);
		}

		public WeightComparer(Func<TWeight, TWeight, int> comparer)
		{
			this.comparer = comparer;
		}
	}
}
